package socket

import (
	"errors"
	"fmt"
	"io"
	"net"
	"strings"
	"time"
)

func SendMessage(protocol, address string, data []byte, timeout time.Duration) ([]byte, error) {
	switch strings.ToLower(protocol) {
	case "tcp":
		return sendTCP(data, address, timeout)
	case "udp":
		return sendUDP(data, address, timeout)
	default:
		return nil, fmt.Errorf("不支持的协议: %s", protocol)
	}
}

// TCP 发送接收
func sendTCP(data []byte, address string, timeout time.Duration) ([]byte, error) {
	conn, err := net.DialTimeout("tcp", address, timeout)
	if err != nil {
		return nil, fmt.Errorf("TCP 连接失败: %w", err)
	}

	err = conn.SetDeadline(time.Now().Add(timeout))
	if err != nil {
		return nil, err
	}

	// === 发送 ===
	_, err = conn.Write(data)
	if err != nil {
		return nil, fmt.Errorf("TCP 发送失败: %w", err)
	}

	if tcpConn, ok := conn.(*net.TCPConn); ok {
		err := tcpConn.CloseWrite()
		if err != nil {
			return nil, err
		}
	} else {
		err := conn.Close()
		if err != nil {
			return nil, err
		}
	}

	// === 接收（循环读取直到 EOF 或超时）===
	buffer := make([]byte, 4096)
	var result []byte
	for {
		n, err := conn.Read(buffer)
		if n > 0 {
			result = append(result, buffer[:n]...)
		}
		if errors.Is(err, io.EOF) {
			// 对端关闭连接，读取完成
			break
		}
		if err != nil {
			return nil, fmt.Errorf("TCP 接收错误: %w", err)
		}
	}

	return result, nil
}

// UDP 发送接收
func sendUDP(data []byte, address string, timeout time.Duration) ([]byte, error) {
	// UDP 是无连接的，使用 WriteTo/ReadFrom
	udpAddr, err := net.ResolveUDPAddr("udp", address)
	if err != nil {
		return nil, fmt.Errorf("解析 UDP 地址失败: %w", err)
	}

	conn, err := net.DialUDP("udp", nil, udpAddr)
	if err != nil {
		return nil, err
	}

	err = conn.SetDeadline(time.Now().Add(timeout))
	if err != nil {
		return nil, err
	}

	// 检查数据报大小（UDP 一般不超过 1400 字节避免分片）
	if len(data) > 1400 {
		return nil, fmt.Errorf("UDP 报文太大（>%d 字节），建议 <1400 字节避免 IP 分片", 1400)
	}

	// === 发送 ===
	_, err = conn.Write(data)
	if err != nil {
		return nil, fmt.Errorf("UDP 发送失败: %w", err)
	}

	// === 接收 ===
	buffer := make([]byte, 4096)
	n, _, err := conn.ReadFromUDP(buffer)
	if err != nil {
		return nil, fmt.Errorf("UDP 接收失败: %w", err)
	}

	return buffer[:n], nil
}
